Applying differentiable programming techniques and machine learning algorithms to foreign programs requires developers to either rewrite their code in a machine learning framework, or otherwise provide derivatives of the foreign code. This paper presents Enzyme, a high-performance automatic differentiation (AD) compiler plugin for the LLVM compiler framework capable of synthesizing gradients of statically analyzable programs expressed in the LLVM intermediate representation (IR). Enzyme synthesizes gradients for programs written in any language whose compiler targets LLVM IR including C, C++, Fortran, Julia, Rust, Swift, MLIR, etc., thereby providing native AD capabilities in these languages. Unlike traditional source-to-source and operator-overloading tools, Enzyme performs AD on optimized IR. On a machine-learning focused benchmark suite including Microsoft's ADBench, AD on optimized IR achieves a geometric mean speedup of 4.2 times over AD on IR before optimization allowing Enzyme to achieve state-of-the-art performance. Packaging Enzyme for PyTorch and TensorFlow provides convenient access to gradients of foreign code with state-of-the-art performance, enabling foreign code to be directly incorporated into existing machine learning workflows.

Summary and Contributions: The key contribution of this paper is a system called Enzyme automatic generation of code for differentiation. While this idea has seen a lot of interest over the last few years, the novelty in this particular proposal is the fact that the generation of code is for the LLVM IR. The main argument for this approach made in the paper is that such generation of code is post-optimization, though intuitively I find it difficult to understand why this is an important feature: it is not obvious that it is a better idea to generate code for computing a derivative before optimization (and to then optimize the generated code using normal compiler tools) than it is to generate code after optimization. That said, the authors do show experimentally that the generate-after-optimization approach is far superior (the generate-before-optimiztation approach is tested as the "ref" option in the paper, and it is often twice as slow as Enzyme). While this non-intuitive result is impressive, I feel that the main argument for the approach is that by doing this at the LLVM level, it is possible to plug the auto-diff software into any LLVM language, there's no more need to work on language-specific auto-diff capabilities.

Applying differentiable programming techniques and machine learning algorithms to foreign programs requires developers to either rewrite their code in a machine learning framework, or otherwise provide derivatives of the foreign code. This paper presents Enzyme, a high-performance automatic differentiation (AD) compiler plugin for the LLVM compiler framework capable of synthesizing gradients of statically analyzable programs expressed in the LLVM intermediate representation (IR). Enzyme synthesizes gradients for programs written in any language whose compiler targets LLVM IR including C, C, Fortran, Julia, Rust, Swift, MLIR, etc., thereby providing native AD capabilities in these languages. Unlike traditional source-to-source and operator-overloading tools, Enzyme performs AD on optimized IR. On a machine-learning focused benchmark suite including Microsoft's ADBench, AD on optimized IR achieves a geometric mean speedup of 4.2 times over AD on IR before optimization allowing Enzyme to achieve state-of-the-art performance.